# 尝试实现异步的训练过程？需要吗？暂时还不需要。
import os 
if os.name == 'nt':
    os.chdir("C:\\federated_malware\\Project")
    print ("cwd",os.getcwd())#获得当前目录
    print ("工作目录",os.path.abspath('.'))#获得当前工作目录
    print ("工作目录",os.path.abspath(os.curdir))#获得当前工作目录
import logging
logger = logging.getLogger()
logger.setLevel(logging.INFO)

# 导入参数------------------------------------------------------
from fl_config import dataset_params,server_params,client_params

logger.info(dataset_params)
logger.info(server_params)
logger.info(client_params)



# 导入模型------------------------------------------------------
from models.minist.dnn_model import Net
# from models.maldroid.dnn_model import Net 
model = Net()
# param_list = model.parameters()

# 导入算法------------------------------------------------------
from trainers.fedbase.FedServerBase import FedServerBase
from trainers.fedbase.FedClientBase import FedClientBase
from trainers.fedavg.FedAvgServer import FedAvgServer
from trainers.fedavg.FedAvgClient import FedAvgClient

# 导入rpc---------------------------------------------------------
from rpc import trainer_service,trainer_proxy
import numpy as np
# ------------------------------------------------fedbase----------------------------------------
# 创建多个客户端和服务器。然后运行base联邦学习。iid数据。客户端和服务器没办法通过重新加载数据进行复用？似乎是没办法。

dataset_name = ['mdm_avg']


base_name = './results/'
algorithm_name = '_fedbase'
suffix = '_.pt'

for name in dataset_name:


    # 为每个客户端生成端口号
    port_list = np.linspace(5000,6000,num=20,dtype=int)
    worker_id_list = []
    for prot in port_list:
        worker_id_list.append("127.0.0.1:{:d}".format(prot))
    
    port_list2 = np.linspace(7000,8000,num=1,dtype=int)
    test_id_list = []
    for prot in port_list2:
        test_id_list.append("127.0.0.1:{:d}".format(prot))
    
    

    train_worker_stub_list = []
    for worker_id in worker_id_list:
        train_worker_stub_list.append(trainer_proxy.FedProxyClient(worker_id))
    
    test_worker_stub_list = []
    for worker_id in test_id_list:
        test_worker_stub_list.append(trainer_proxy.FedProxyClient(worker_id))
    server = FedServerBase((train_worker_stub_list,test_worker_stub_list),model,server_params,client_params)
        
    logger.info("本次测试的数据集的名字:{}".format(name))
    logger.info("本次测试的算法的名字:{}".format(algorithm_name))
    server.start_federated()
    server.save_state(base_name+name+algorithm_name+suffix)




# ------------------------------------------------fedavg----------------------------------------
# 创建多个客户端和服务器。然后运行base联邦学习。iid数据。客户端和服务器没办法通过重新加载数据进行复用？似乎是没办法。

# train_dataset_list = [mdm_avg,mdm_ln,mdm_ld,mdm_fn,mdm_quantity]
# test_dataset_list = [mdm_test,mdm_test,mdm_test,mdm_test,mdm_test]
# dataset_name = ['mdm_avg','mdm_ln','mdm_ld','mdm_fn','mdm_quantity']
# train_dataset_list = [fdm_noniid]
# test_dataset_list = [fdm_test]
# dataset_name = ['fdm_noniid']

# train_dataset_list = [maldm_avg,maldm_ln,maldm_ld,maldm_fn,maldm_quantity]
# test_dataset_list = [maldm_test,maldm_test,maldm_test,maldm_test,maldm_test]
# dataset_name = ['maldm_avg','maldm_ln','maldm_ld','maldm_fn','maldm_quantity']

# base_name = './results/'
# algorithm_name = '_fedavg'
# suffix = '_.pt'

# for name,train_data,test_dataset in zip(dataset_name,train_dataset_list,test_dataset_list):
#     train_worker_list = [FedAvgClient(train_dataset,model) for train_dataset in train_data]
#     test_worker_list = [FedAvgClient(test_dataset,model) for test_dataset in test_dataset]
#     server = FedAvgServer((train_worker_list,test_worker_list),model,server_params,client_params)
    
    
#     logger.info("本次测试的数据集的名字:{}".format(name))
#     logger.info("本次测试的算法的名字:{}".format(algorithm_name))
#     server.start_federated()
#     server.save_state(base_name+name+algorithm_name+suffix)

